!Subject: snchi.f
* Q+D to evaluate chi^2 for an externally provided (O_m, O_l, w)
* Derived from frwfit used for Tonry et al. 2003, ApJ, 594, xxx
* 030807 - rev 1.01  (John Tonry)
*
* Slightly modified by Max Tegmark August 2003 to 
* not give NaN for cases where there's no big bang,
* and to auto-initialize if needed.

      program test
      call sn1ademo
      end
      
      
      subroutine sn1ademo
      implicit none
      real*8 Om, Ol, w, wa, chi2
      Om = 0.3
      Ol = 0.8
      w  = -0.9 
      wa = 0.0
      !Om = 0.3
      !Ol = 1.7
      !w  = -1
      call compute_sn1a_chi2(Om,Ol,w,wa,chi2)
      !write(6,'(5f10.3)') Om, Ol, w, wa,chi2
      ! Test against this:
      ! snchi.x om=0.3 ov=0.8 w=-0.9
      return
      end

      subroutine compute_sn1a_chi2(Om,Ol,w,wa,chi2)
      ! Om = Omega_m
      ! Ol = Omega_l
      ! w  = Equation of state w
      ! wa defined as: w = w + (1-a)*wa
      !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
      !
      implicit none
      real*8 Om, Ol, w, wa, chi2
        integer npt, i, failure
      integer MAXPT
      parameter (MAXPT=500)
      real*8 z(MAXPT), dist(MAXPT), ddist(MAXPT), dpred(MAXPT)
      real*8 systematic, Ok, a, tmp, epsilon
      real*8 Or, have, hwgt, time, dmet, dprop
      logical initialized
      save initialized, z, dist, ddist
      data initialized/.false./

      if (.not.initialized) then
        print *,'Initializing SN Ia likelihood stuff...'
        call load_sn1a_data_riess(npt,z,dist,ddist)
        initialized = .true.
      end if

* Defaults for parameters
      Or = 0
      systematic = 0
      
      !!!!!!!!!!!!!!!!!!!
      ! Added by Max:
      ! Test if we're in the forbidden zone where the integrand blows up
      ! (where's H^2 < 0 for some a between 0 and 1), and if so returns
      ! chi2 = 666 instead of chi2 = NaN. This code works only for Or=0.
      ! The function f(a) = Omega_m + Omega_k*a + Omega_l*a**(-3*w)
      ! can have at most one local minimum, at (Ok/(3*w*Ol))**(-1/(1+3*w)), 
      ! so simply check if f(a)<0 there or at the endpoints a=0, a=1.
      ! g(0) = Omega_m - Omega_l*a**(-3*w) < 0 if w > 0 & Omega_k > 1
      !                                     or if w = 0 & Omega_l < 1       
      ! g(1) = Omega_m + Omega_k + Omega_l = 1 > 0
      failure = 0
      epsilon = 0
      !epsilon = 0.04  ! Numerical integration fails even before H^2 goes negative.
      Ok = 1-Om-Ol ! Spatial curvature
      if (w*Ol.ne.0) then
        tmp = Ok/(3*w*Ol)
        if ((tmp.gt.0).and.(1+3*w.ne.0)) then ! f'(0)=0 for some a>0
          a = tmp**(-1/(1+3*w))
          if (a.lt.1) then
            if (Om + Ok*a + Ol*a**(-3*w).lt.epsilon) failure = 1
          end if
        end if
      end if
      if ((w.eq.0).and.(Ok.gt.1)) failure = 2
      if ((w.gt.0).and.(Ol.lt.0)) failure = 3
      if (failure.gt.0) then
        print *,'Big Bang failure mode ',failure
        chi2 = 666
        return
      end if
      !!!!!!!!!!!!!!!!!!!


* Calculate the cosmology for these parameters
      have = 0
      hwgt = 0
      do 20 i = 1,npt
         call friedmann(z(i),Om,Ol,w,wa,Or,dmet,dprop,time)
         !write(*,'(9f12.5)') z(i),Om,Ol,w,wa,Or,dmet,dprop,time
         if (dprop.le.0) then ! Mathematically impossible, but still happens for high z
           print *,'Failure mode 4' 
           chi2 = 666
           return
         end if
         if (.not.((dprop.ge.0).or.(dprop.le.0))) then
           ! Sometimes the integral bombs numerically even when it's mathematically sound.
           failure = 5  ! dprop = NaN
           chi2 = 666
           print *,'Failure mode 5'
           return
         end if
         if (dprop.gt.1e30) then ! dprop fapp infinite
           print *,'Failure mode 7' 
           chi2 = 666
           return
         end if
         dpred(i) = log10(dprop*(1+z(i))) + 
     $        0.2*systematic/log10(2.0)*log10(1+z(i))
         hwgt = hwgt + 1/ddist(i)**2
         have = have + (dist(i)-dpred(i))/ddist(i)**2
         !print *,i,z(i),dprop,dpred(i),hwgt,have
 20   continue

*     Sum up Chi^2
!     The only stuff from above that's used here is 
!     have, hwgt, dpred(i), dist(i), ddist(i)
* Marginalize over H0: since H0 appears linearly in chi^2's arg, the
* marginalization over it simply means evaluate the weighted average and
* evaluate at that H0.  Actual integration over a ln(H0) would also 
* put in a factor of sqrt(2pi/hwgt), and the prior for ln(H0) might be
* something like ln(10), so the Bayesian marginalization would be this
* result, but with a factor of sqrt(2pi/hwgt)/ln(10) as well.  Since this
* is the same for all (Om,Ol) that I try, and since I'm normalizing the
* probability anyway, I'll just ignore this factor.
      have = have / hwgt
      chi2 = 0
      do 23 i = 1,npt
        !write(*,'(i4,9f12.5)') i, z(i), dist(i), dpred(i), ddist(i), have
        chi2 = chi2 + ((dist(i)-dpred(i)-have)/ddist(i))**2
 23   continue

      !prob = exp(-0.5*(chi2-npt))
      if (.not.((chi2.ge.0).or.(chi2.le.0))) then
        ! Shouldn't happen, but check just in case.
        failure = 6  ! chi2 = NaN
        print *,'Failure mode 6',have,chi2
        chi2 = 666
      end if

      write(*,'(i4,5f10.3)') npt, Om, Ol, w, wa, chi2

      return
      end

      subroutine friedmann(z,Om,Ol,w,wa,Or,dist,dprop,time)
      ! Max comment: omega = (Om,Or,Ol), where Or = Omega_radiation ~ 0 
      ! time = t from z to now
      ! dist = R0*r1
      ! dprop = R0 * Sk(r1)
      implicit none
      real*8 z, Om, Ol, w, wa,Or, dist, dprop, time
      real*8 omega(3), rgral, curve, a0, a1, clever
      !real*8 tgral
      external rgral
      !external tgral
      omega(1) = Om
      omega(2) = Or
      omega(3) = Ol
      curve = 1 - Om - Ol - Or
      a0 = 1/(1+z)
      a1 = 1
      !print *,omega,w,wa,a0,a1,rgral(omega,w,wa,a0)
      dist = clever(omega,w,wa,a0,a1,rgral)
      !time = clever(omega,w,a0,a1,tgral)
      time = 666 ! Save time by not computing, since never used
      if(curve .lt. 0) then
         dprop = sin(sqrt(abs(curve))*dist) / sqrt(abs(curve))
      else if(curve .gt. 0) then
         dprop = sinh(sqrt(abs(curve))*dist) / sqrt(abs(curve))
      else
         dprop = dist
      end if
      return
      end

      real*8 function rgral(omega, w, wa, a)
      implicit none
      real*8 omega(3), w, wa, a, curve
      curve =  1 - omega(1) - omega(2) - omega(3)
      ! Max comment: rgral**(-2) = a^4 (H/H0)^2 = 
      ! = a^4 (Ok/a^2 + Om/a^3 + Or/a^4 + Ol/a^(3*(1+w)))
      ! = Ok a^2 + Om a + Or + Ol^(-3*(1+w)+4)
      ! omega = (Om,Or,Ol), where Or = Omega_radiation ~ 0
      rgral = 1 / sqrt(curve*a*a + omega(1)*a+omega(2)+
     $   omega(3)*exp(-3.0*((1.0 + w + wa) * log(a) - wa*a + wa)))
      !  omega(3)*a**(-3*(1+w)+4))
      return
      end

      real*8 function tgral(omega, w, a)
      implicit none
      real*8 omega(3), w, a, curve
      curve =  1 - omega(1) - omega(2) - omega(3)
      tgral = 1 / sqrt(curve + omega(1)/a + omega(2)/a/a +
     $   omega(3)*a**(-3*(1+w)+2))
      return
      end

      REAL*8 FUNCTION CLEVER(S,T,T2,A,B,F)
C
C     CLEVER IS A SMALL MODIFICATION OF SIMDON, WRITTEN BY
C     FRANK NAGY AT CALTECH.
C
C     THIS ADAPTIVE SIMPSON RULE IS FROM  NUMERICAL COMPUTING  BY
C     SHAMPINE AND ALLEN.
C
C     F      = FUNCTION TO BE INTEGRATED
C     S,T,T2    = PARAMETERS PASSED TO INTEGRAND FUNCTION
C     A,B    = LOWER AND UPPER LIMITS OF INTEGRATION
C     CLEVER = RESULT OF THE INTEGRATION
C     AREA   = RESULT OF INTEGRATING ABS(F) FROM A  TO  B
C     ERROR  = ESTIMATED ERROR OF ANS
C     ACC    = DESIRED ACCURACY (ERROR/AREA)
C     IFLAG  = 1 FOR NORMAL RETURN
C            = 2 IF IT IS NECESSARY TO GO TO 200 SUBLEVELS OR
C            = 4 IF SUBINTERVALS GET TOO SMALL FOR MACHINE WORD LENGTH
C            = 3 IF MORE THAN 2000 FUNCTION EVALUATIONS USED. ROUGH
C                APPROXIMATIONS ARE USED TO COMPLETE THE COMPUTATIONS.
C
      implicit none
      integer LORR(200), IFLAG, LVL
      real*8 FV(5),FIT(200),F2T(200),F3T(200),DAT(200),
     1 ARESTT(200),ESTT(200),EPST(200),PSUM(200), F
      real*8 U, ACC, FOURU, EPS, ERROR, ALPHA, DA, AREA, AREST, KOUNT
      real*8 WT, EST, DX, ESTL, ESTR, SUM
      real*8 ARESTL, ARESTR, DIFF, A, B, S, T, T2
      EXTERNAL F
C
C     SET U TO APPROXIMATELY UNIT ROUND-OFF OF THE IBM 360/67
C     U IS TWICE AS GOOD ON CDC 7600, SO TAKE; (8/22/73)
C
      U=1.0E-7
C
C     INITIALIZE
C
      ACC=1.0E-4
      FOURU=4.0*U
      IFLAG=1
      EPS=ACC
      ERROR=0.0
      LVL=1
      LORR(LVL)=1
      PSUM(LVL)=0.0
      ALPHA=A
      DA=B-A
      AREA=0.0
      AREST=0.0
      FV(1)=F(S,T,T2,ALPHA)
      FV(3)=F(S,T,T2,ALPHA+0.5*DA)
      FV(5)=F(S,T,T2,ALPHA+DA)
      KOUNT=3
      WT=DA/6.
      EST=WT*(FV(1)+4.*FV(3)+FV(5))
C
C     BASIC STEP . HAVE ESTIMATE EST OF INTEGRAL ON (ALPHA,ALPHA+DA).
C     BISECT AND COMPUTE ESTIMATES ON LEFT AND RIGHT HALF INTERVALS.
C     SIMILARLY TREAT INTEGRAL OF ABS(F(X)). SUM IS BETTER VALUE FOR
C     INTEGRAL AND DIFF/15. IS APPROXIMATELY ITS ERROR
C
 1    DX=.5*DA
      FV(2)=F(S,T,T2,ALPHA+0.5*DX)
      FV(4)=F(S,T,T2,ALPHA+1.5*DX)
      KOUNT=KOUNT+2
      WT=DX/6.
      ESTL=WT*(FV(1)+4.*FV(2)+FV(3))
      ESTR=WT*(FV(3)+4.*FV(4)+FV(5))
      SUM=ESTL+ESTR
      ! Added by Max:
      if (.not.((SUM.ge.0).or.(SUM.le.0))) return ! Sum = NaN, so might as well bail out right now to save time      
      ARESTL=WT*(ABS(FV(1))+ABS(4.*FV(2))+ABS(FV(3)))
      ARESTR=WT*(ABS(FV(3))+ABS(4.*FV(4))+ABS(FV(5)))
      AREA=AREA+((ARESTL+ARESTR)-AREST)
      DIFF=EST-SUM
C
C     IF ERROR UNACCEPTABLE, GO TO 2. IF INTERVAL TOO SMALL OR
C     TOO MANY LEVELS OR TOO MANY FUNCTION EVALUATIONS, SET A FLAG
C     AND GO TO 2 ANYWAY.
C
      IF(ABS(DIFF).LE.EPS*ABS(AREA)) GO TO 2
      IF(ABS(DX).LE.FOURU*ABS(ALPHA)) GO TO 51
      IF(LVL.GE.200) GO TO 5
      IF(KOUNT.GE.2000) GO TO 6
C
C     HERE TO RAISE LEVEL. STORE INFORMATION TO PROCESS RIGHT HALF
C     INTERVAL LATER. INITIALIZE FOR  BASIC STEP  SO AS TO TREAT
C     LEFT HALF INTERVAL.
C
      LVL=LVL+1
      LORR(LVL)=0
      FIT(LVL)=FV(3)
      F2T(LVL)=FV(4)
      F3T(LVL)=FV(5)
      DA=DX
      DAT(LVL)=DX
      AREST=ARESTL
      ARESTT(LVL)=ARESTR
      EST=ESTL
      ESTT(LVL)=ESTR
      EPS=EPS/1.4
      EPST(LVL)=EPS
      FV(5)=FV(3)
      FV(3)=FV(2)
      GO TO 1
C
C     ACCEPT APPROXIMATE INTEGRAL SUM. IF IT WAS ON A LEFT INTERVAL,
C     GO TO  MOVE RIGHT . IF A RIGHT INTERVAL, ADD RESULTS TO FINISH
C     AT THIS LEVEL. ARRAY LORR(MNEMONIC FOR LEFT OR RIGHT)
C     TELLS WHETHER LEFT OF RIGHT INTERVAL AT EACH LEVEL.
C
 2    ERROR=ERROR+DIFF/15.
 3    IF(LORR(LVL).EQ.0) GO TO 4
      SUM=PSUM(LVL)+SUM
      LVL=LVL-1
      IF(LVL.GT.1) GO TO 3
      CLEVER=SUM
      RETURN
C
C     MOVE RIGHT . RESTORE SAVED INFORMATION TO PROCESS RIGHT HALF
C     INTERVAL.
C
 4    PSUM(LVL)=SUM
      LORR(LVL)=1
      ALPHA=ALPHA+DA
      DA=DAT(LVL)
      FV(1)=FIT(LVL)
      FV(3)=F2T(LVL)
      FV(5)=F3T(LVL)
      AREST=ARESTT(LVL)
      EST=ESTT(LVL)
      EPS=EPST(LVL)
      GO TO 1
C
C     ACCEPT  POOR  VALUE. SET APPROPRIATE FLAGS.
C
 5    IFLAG=2
      GO TO 2
 6    IFLAG=3
      GO TO 2
 51   IFLAG=4
      GO TO 2
      END

      real*8 function log10(x)
      implicit none
      real*8 x
      log10 = log(x)/log(10.d0)
      return
      end
      
      subroutine load_sn1a_data(npt,z,dist,ddist)
      !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
      ! This subroutine is a trivial adaption by Max Tegmark of the code snchi.f
      ! from John Tonry; see astro-ph/0305???.
      ! From jt@avidya.ifa.hawaii.edu  Fri Aug  8 01:58:45 2003
      ! Derived from frwfit used for Tonry et al. 2003, ApJ, 594, xxx
      ! 030807 - rev 1.01  (John Tonry)
      !
      implicit none
      integer IREFMIN, IREFMAX, npt, i, ii, lnblnk
      integer MAXPT
      parameter (MAXPT=500)
      real*8 l2(MAXPT), b2(MAXPT), z(MAXPT), dist(MAXPT)
      real*8 ddist(MAXPT), av(MAXPT)
      real*8 c, ZMIN, ZMAX, AVCUT, VRANDOM, czlog
      integer ncontrib(MAXPT), iref(MAXPT)
      character*4 sname(MAXPT)
      character*5 contrib(MAXPT), exclude
      character*8 ref(MAXPT)
      character*10 host(MAXPT)
      character*80 datafile
      print *,'Loading SN Ia data...'

      C = 299792.458

* Cuts on the data
      ZMIN = 0.01
* ZMAX = 1.3 cuts out sn97ff
      ZMAX = 2.0
C       ZMAX = 1.3
      AVCUT = 0.50
* Cut by analysis list
      EXCLUDE = 'none '
C       EXCLUDE = 'P    '
* Minimum reference (15 cuts Riess98, SCP, Riess99, etc)
      IREFMIN = -1
C       IREFMIN = 15

* Maximum reference (41 cuts Germany)
      IREFMAX = 1000
C       IREFMAX = 41

* What is the uncertainty in redshift and/or LSF amplitude?
      VRANDOM = 500

      datafile = 'sntable.txt'

* Suck down the data
      open(unit=2,file=datafile, status='old')
      read(2,*)
      npt = 0
      do 5 i = 1,MAXPT
         read(2,1000,end=6) sname(npt+1), l2(npt+1), b2(npt+1), 
     $        host(npt+1), ref(npt+1), czlog,
     $        dist(npt+1), ddist(npt+1), av(npt+1), contrib(npt+1),
     $        ncontrib(npt+1)
 1000    format(2x,a,2f8.3,1x,2a,3f7.3,f6.2,1x,a,i3,3f7.3,f7.2)

         ii = index(ref(npt+1),':') - 1
         if(ii.le.0) ii = len(ref(npt+1))
         read(ref(npt+1)(:ii-1),'(i2)') iref(npt+1)

* Apply any cuts
         if(czlog .le. log10(C*ZMIN)) goto 5
         if(czlog .ge. log10(C*ZMAX)) goto 5
         if(av(npt+1) .ge. AVCUT) goto 5
         if(contrib(npt+1) .eq. EXCLUDE) goto 5
         if(czlog .ge. log10(c*0.1) .and. 
     $        iref(npt+1).lt.IREFMIN) goto 5
         if(iref(npt+1).gt.IREFMAX) goto 5

* Acceptable point
         npt = npt + 1
         if (npt.gt.MAXPT) stop 'SN Ia DEATH ERROR: MAXPT too small'
* Convert log(cz) to z
         z(npt) = 10.**czlog / C
* Convert log(c*dist) to log(dist)
         dist(npt) = dist(npt) - log10(C)
* Augment ddist by velocity uncertainty
         ddist(npt) = sqrt(ddist(npt)**2 + 
     $        (VRANDOM/(C*z(npt)*alog(10.)))**2)

5     continue
6     close(2)
      print *,npt,' supernovae loaded from '//datafile(1:lnblnk(datafile))
      return
      end

      subroutine load_sn1a_data_riess(npt,z,dist,ddist)
      !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
      ! This subroutine is a trivial adaption by Max Tegmark of the code snchi.f
      ! from John Tonry; see astro-ph/0305???.
      ! From jt@avidya.ifa.hawaii.edu  Fri Aug  8 01:58:45 2003
      ! Derived from frwfit used for Tonry et al. 2003, ApJ, 594, xxx
      ! 030807 - rev 1.01  (John Tonry)
      !
      implicit none
      integer npt, i, lnblnk
      integer MAXPT
      parameter (MAXPT=500)
      real*8 z(MAXPT), dist(MAXPT)
      real*8 ddist(MAXPT)
      real*8 c, VRANDOM
      character*20 sname(MAXPT), stat(MAXPT)
      character*80 datafile
      print *,'Loading SN Ia data...'

      C = 299792.458

* What is the uncertainty in redshift and/or LSF amplitude?
      VRANDOM = 500

      datafile = '/home/mjarvis/code/corr/riess.txt'

* Suck down the data
      open(unit=2,file=datafile, status='old')
      read(2,*)
      npt = 0
      do 5 i = 1,MAXPT
         read(2,*,end=6) sname(npt+1), z,
     $        dist(npt+1), ddist(npt+1), stat(npt+1)

* Apply any cuts
        if (stat(npt+1) .ne. 'gold') goto 5

* Acceptable point
         npt = npt + 1
         if (npt.gt.MAXPT) stop 'SN Ia DEATH ERROR: MAXPT too small'
* Convert dist_modulus to log(dist)
         dist(npt) = dist(npt)/5.-5.
* Augment ddist by velocity uncertainty
         ddist(npt) = sqrt(ddist(npt)**2 + 
     $        (VRANDOM/(C*z(npt)*alog(10.)))**2)

5     continue
6     close(2)
      print *,npt,' supernovae loaded from '//datafile(1:lnblnk(datafile))
      return
      end
